from transfer import instantiate_from_config
import argparse
import pytorch_lightning as pl
import yaml
import os
from pytorch_lightning.callbacks import ModelCheckpoint
import pytorch_lightning.loggers as pl_loggers
import torch
import cv2
from pytorch_lightning.utilities import rank_zero_only

os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"


def parser():
    args_parser = argparse.ArgumentParser(description="TransferBasedDetection")
    args_parser.add_argument("--config", type=str, required=True)
    args_parser.add_argument("--train", action="store_true", default=False)
    args_parser.add_argument("--test", action="store_true", default=False)
    args_parser.add_argument("--predict", action="store_true", default=False)
    return args_parser.parse_args()


class ImageLogger(pl.Callback):
    def __init__(self, root_dir, log_step_interval):
        self.root_dir = root_dir
        self.log_step_interval = log_step_interval

    @rank_zero_only
    def on_validation_start(self, trainer, pl_module: pl.LightningModule) -> None:
        if not os.path.exists(os.path.join(self.root_dir, "images")):
            os.makedirs(os.path.join(self.root_dir, "images"))

    @rank_zero_only
    def on_test_start(self, trainer, pl_module: pl.LightningModule) -> None:
        if not os.path.exists(os.path.join(self.root_dir, "images")):
            os.makedirs(os.path.join(self.root_dir, "images"))

    @rank_zero_only
    def on_validation_batch_end(
        self, trainer, pl_module: pl.LightningModule, outputs, batch, batch_idx
    ) -> None:
        if batch_idx % self.log_step_interval == 0:
            log = pl_module.log_images(batch, prefix="val")
            for k, v in log.items():
                assert v.dtype == torch.uint8
                pl_module.logger.experiment.add_image(k, v, pl_module.global_step)
                path = os.path.join(
                    self.root_dir,
                    "images",
                    f"{k}_{pl_module.global_step}_{batch_idx}.png",
                )
                if not os.path.exists(os.path.dirname(path)):
                    os.makedirs(os.path.dirname(path))
                cv2.imwrite(path, v.permute(1, 2, 0).cpu().numpy())

    @rank_zero_only
    def on_test_batch_end(
        self, trainer, pl_module: pl.LightningModule, outputs, batch, batch_idx
    ) -> None:
        assert isinstance(outputs[0], dict)
        for k, v in outputs.items():
            pl_module.logger.experiment.add_image(k, v, batch_idx)
            path = os.path.join(self.root_dir, "images", f"{k}_{batch_idx}.png")
            if not os.path.exists(os.path.dirname(path)):
                os.makedirs(os.path.dirname(path))
            cv2.imwrite(path, v.permute(1, 2, 0).cpu().numpy())

    @rank_zero_only
    def on_predict_batch_end(
        self, trainer, pl_module: pl.LightningModule, outputs, batch, batch_idx
    ) -> None:
        assert isinstance(outputs[0], dict)
        for k, v in outputs[0].items():
            pl_module.logger.experiment.add_image(k, v, batch_idx)
            path = os.path.join(self.root_dir, "images", f"{k}_{batch_idx}.png")
            if not os.path.exists(os.path.dirname(path)):
                os.makedirs(os.path.dirname(path))
            cv2.imwrite(path, v.permute(1, 2, 0).cpu().numpy())


def _TransferBasedDetection(name: str, config: dict):
    ## root_dir ##
    root_dir = os.path.join(".", "logs", name)
    ## model ##
    model_config = config.pop("model")
    monitor = model_config.pop("monitor")
    model = instantiate_from_config(model_config)
    ## prepare data ##
    data_module = instantiate_from_config(config.pop("data"))
    ## loggers ##
    tb = pl_loggers.TensorBoardLogger(root_dir, name=None, purge_step=5)
    image_logger = ImageLogger(tb.log_dir, **config.pop("logger"))
    ## checkpointing ##
    checkpoint_callback = ModelCheckpoint(
        dirpath=os.path.join(tb.log_dir, "models"),
        filename="{epoch:03d}-{val/loss:.6f}",
        save_top_k=3,
        monitor=monitor,
        save_last=True,
        auto_insert_metric_name=False,
        verbose=True,
    )
    ## trainer ##
    trainer = pl.Trainer(
        default_root_dir=root_dir,
        logger=tb,
        callbacks=[checkpoint_callback, image_logger],
        accelerator="gpu",
        devices=-1,
        strategy="ddp_find_unused_parameters_true",
        **config.pop("trainer"),
    )
    return trainer, model, data_module


def main(args):
    # 仅支持train，test，predict三种模式中的一种
    assert (
        args.train + args.test + args.predict == 1
    ), "Please choose one mode from train, test, predict"
    with open(args.config, "r") as f:
        config = yaml.load(f, Loader=yaml.FullLoader)
    name = os.path.basename(args.config).split(".")[0]
    trainer, model, data_module = _TransferBasedDetection(name, config)
    torch.set_float32_matmul_precision("high")
    if args.train:
        trainer.fit(model=model, datamodule=data_module)
    if args.test:
        trainer.test(model=model, datamodule=data_module)
    if args.predict:
        trainer.predict(model=model, datamodule=data_module)


if __name__ == "__main__":
    args = parser()
    main(args)
